from rag_system import LocalRAGSystem
import os

# 设置硅基流动API密钥
os.environ["GUIJI_API_KEY"] = "your-guiji-api-key-here"  # 替换为您的API密钥

# 初始化RAG系统
rag = LocalRAGSystem()

# 加载文档
rag.load_documents("./documents")  # 确保documents目录中有.txt文件

# 进行问答
chat_history = []
while True:
    query = input("\n请输入您的问题 (输入'quit'退出): ")
    if query.lower() == 'quit':
        break
        
    result = rag.query(query, chat_history)
    
    print("\n回答:", result["answer"])
    print("\n参考文档:")
    for doc in result["source_documents"]:
        print("- 文档片段:", doc.page_content[:200], "...")
        print("  来源:", doc.metadata)
        print()
    
    # 更新对话历史
    chat_history.append((query, result["answer"])) 